# @Author:黄景涛
# @Date: 2023/10/9

import time

import re
import random

import requests
from collections import deque
from multiprocessing import Process, Queue, Lock
from concurrent.futures import ThreadPoolExecutor
from pypinyin import lazy_pinyin

from settings import fk, create_organization, Node, dfs
from user_manage_login import login


class OrganizateError(Exception):
    """组织校验的错误"""


class UserError(Exception):
    """用户校验的错误"""


# 线程数量控制
THREAD_LIMIT = 30

# 请求头配置
headers = {
    "User-Agent": fk.user_agent(),
    "Cookie": "",
    "X-Csrftoken": ""
}


def update_headers():
    cookies, token = login()
    # 更新headers中Cookie键值对, X-Csrftoken
    headers['Cookie'] = cookies
    headers['X-Csrftoken'] = token


# 域名
base_url = "http://paas.test.com"

# full_url
full_organizate_url = base_url + '/o/bk_user_manage/api/v1/web/departments/'
full_user_url = base_url + '/o/bk_user_manage/api/v1/web/profiles/'


# 全局的session对象
session = requests.session()


def _parser_id(res_data: dict) -> int:
    """
    解析响应结果中，组织id
    :param res_data: 请求的响应json
    :return: 组织id
    """
    pattern = "'id': (.*?),"
    organizate_id = re.search(pattern, str(res_data)).group(1)
    return int(organizate_id)


def _update_organizate_structure(target_node: Node, organizate_id: int):
    """
    更新节点的属性id
    :param target_node: 目标组织
    :param organizate_id: 组织id
    :return: None
    """
    target_node.id = organizate_id


def add_organizate(target_node: Node, name: str, headers, category_id: int = 1, parent=None, retry_times=3):
    """
    1.发起异步请求：新增组织； 2.更新组织字典（往原组织名中，补充组织id，用于后续的用户添加）
    :param target_node: 目标组织
    :param name: 组织名称
    :param headers: 请求头，这里本来可以不传，但是在多线程执行时，必须以形参传入，否则没法拿到全局headers的最新值
    :param category_id: 类型
    :param parent: 父级组织id
    :param retry_times: 请求重试次数
    :return: None
    """
    # 参数判断
    data = {
        "name": name,
        "category_id": category_id,
    }

    if parent is not None:
        data.update(parent=parent)
    # print(f'{name}的参数：{data}')

    i = 0
    while True:
        if i > retry_times - 1:
            break
        try:
            resp = session.post(full_organizate_url, json=data, headers=headers)
            res_data = resp.json()
            # print(f'新增组织{name}响应结果：{res_data}')
            if res_data.get('code') == -1:
                msg = f'{name}新增失败 >>{res_data.get("message")}'
                raise OrganizateError(msg)
        # except requests.exceptions.JSONDecodeError:
        #     # TODO
        #     """ 捕获该类异常，表示登录过期，执行登录更新了Cookie和token后，再继续进程"""
        #     break
        except Exception as e:
            msg = f'错误：{str(e)[:100]}'
            print(f'organizate：{name} =>> 失败：{msg}')
            i += 1
        else:
            # 提取响应中组织id
            organizate_id = _parser_id(res_data)
            # 更新目标字典中，组织的信息(添加组织id)
            _update_organizate_structure(target_node, organizate_id)
            print(f'organizate：{name} =>> 成功')
            return True


def add_user(success_set: set, target_node: Node, headers, category_id: int = 1, retry_times=3):
    """
    新增用户
    :param success_set: 用户名的集合，用于过滤重复的用户名
    :param target_node: 目标组织
    :param headers: 请求头，这里本来可以不传，但是在多线程执行时，必须以形参传入，否则没法拿到全局headers的最新值
    :param category_id: 组织类型
    :param retry_times: 请求重试次数
    :return: None
    """
    # 参数构建
    # 用户名规则：由1-32位字母、数字、下划线_、点.、减号-组成，已字母或数字开头
    character_pool = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz.-_"
    # 中文名：随机姓名+后缀
    display_name = fk.name() + ''.join(random.choices(character_pool, k=random.randint(0, 5)))
    # 姓名拼音+后缀
    pinyin = ''.join(lazy_pinyin(display_name))
    # print(pinyin)
    while True:
        username = pinyin + ''.join(random.choices(character_pool, k=random.randint(2, 31 - len(pinyin))))
        if username not in success_set:
            success_set.add(username)
            break
    # 性别：0, 男；1，女
    sex_choice = (0, 1)
    # 账户状态，大概率取NORMAL：正常，小概率取其他
    status_choice = ("NORMAL", "LOCKED", "DELETED", "DISABLED", "EXPIRED")
    # 在职状态：大概率取IN: 在职，小概率取离职
    staff_status_choice = ("IN", "OUT")

    data = {
        "category_id": category_id,
        "leader": [],
        "departments": [target_node.id],  # 组织
        "password_valid_days": -1,  # 密码有效期，-1表示永久
        "username": username,  # 用户名，唯一
        "display_name": display_name,  # 用户全名，可以相同
        "email": fk.free_email(),  # 邮箱
        "telephone": fk.phone_number(),  # 手机号
        "iso_code": "cn",
        "status": ''.join(random.choices(status_choice, (0.9, 0.025, 0.025, 0.025, 0.025), k=1)),  # 账户状态
        "staff_status": ''.join(random.choices(staff_status_choice, (0.9, 0.1), k=1)),  # 在职状态
        "position": "0",
        "wx_userid": random.choice(("", "weixin" + ''.join(random.choices(character_pool, k=10)))),  # 微信
        "qq": random.choice(("", ''.join(random.choices(character_pool[:10], k=random.randint(6, 12))))),  # qq
        "account_expiration_date": "",  # 账号到期时间
        "sex": random.choice(sex_choice)  # 性别
    }

    i = 0
    while True:
        if i > retry_times - 1:
            break
        try:
            resp = session.post(full_user_url, json=data, headers=headers)
            res_data = resp.json()
            if res_data.get('code') == -1:
                msg = f'{username}新增失败，错误：{res_data.get("message")}'
                raise UserError(msg)
        # except requests.exceptions.JSONDecodeError:
        #     """ 捕获该类异常，表示登录过期，执行登录更新了Cookie和token后，再继续进程"""
        #     break
        except Exception as e:
            msg = f'{username}新增失败，错误：{str(e)[:100]}'
            print(f'user：{display_name} ->> 失败{i}:{msg}')
            i += 1
        else:
            print(f'user：{display_name} ->> 成功')
            return True


# 子任务：创建每个组织
def organizate_subtask(current_node: Node, headers):
    """
    创建组织
    :param current_node: 组织节点
    :param headers: 请求头
    :return: None
    """
    # root节点，跳过创建组织的请求
    if current_node.name == 'root':
        return

    # 1.新增组织的请求
    parent_node = current_node.parent
    if parent_node.name == 'root':
        success_flag = add_organizate(current_node, current_node.name, headers)
    else:
        # print(parent_node.id, parent_node.name)
        success_flag = add_organizate(current_node, current_node.name, headers, parent=parent_node.id)
    # 请求结果=True，返回创建成功的组织节点
    if success_flag:
        return current_node


# 任务1：创建组织，并把节点对象，通过管道传给任务2
def organizate_task(q: Queue, root: Node, headers):
    """
    层序遍历（BFS）：每遍历一层级组织，一边创建组织，一边创建组织下的用户
    :param root: 组织架构树的根节点，虚拟节点
    :param headers: 请求头
    :return: None
    """
    # BFS，层序遍历
    queue = deque()
    queue.append(root)

    with ThreadPoolExecutor(THREAD_LIMIT) as tpool:
        while queue:
            # 从队列左侧出队，一个节点
            current_node = queue.popleft()
            """节点的处理操作"""
            # 向线程池中，提交任务
            target_node = tpool.submit(organizate_subtask, current_node, headers).result()
            # 把任务返回值，塞进数据管道；过滤：root节点/创建失败的节点组织，不处理
            if target_node:
                q.put(target_node)
            # 将当前节点的子节点，入队
            for child in current_node.childrens:
                queue.append(child)
        # q管道中，消息消费殆尽，传一个退出的信号量
        q.put(-1)
        tpool.shutdown()


# 任务2：在目标节点上，创建用户
def user_task(q: Queue, success_set: set, headers):
    with ThreadPoolExecutor(THREAD_LIMIT) as tpool:
        while True:
            # 从管道队列中，消费消息
            current_node = q.get()
            if current_node == -1:  # 接收到信号量时，退出while循环
                break

            user_limit = current_node.user_limit
            for i in range(user_limit):
                tpool.submit(add_user, success_set, current_node, headers)


def engine(root: Node):
    """
    主函数入口
    :param root: 虚拟的根节点
    :return: None
    """
    update_headers()

    # 构造用户名的过滤器，防止用户唯一标识重复
    success_set = set()
    # 创建管道队列
    q = Queue()

    # 创建2个进程对象
    p1 = Process(target=organizate_task, args=(q, root, headers))
    p2 = Process(target=user_task, args=(q, success_set, headers))

    # 启动进程
    p1.start()
    p2.start()


if __name__ == '__main__':
    root = create_organization()
    # 启动协程
    engine(root)
